feat(attention): fixed-m splash attention kernel with dynamic bounds and safety fallbacks - #477

…and safety fallbacks

Implements exact fixed-m splash attention in Pallas on TPU. Instead of an
online running max, a head / Q-block whose Cauchy-Schwarz bound
max||q|| * max||k - k_bar|| is within a length-dependent window uses a
precomputed m = ceil(bound) - C(N) (+ q . k_bar with virtual K-centering) and
skips the per-block rescales; everything else keeps online softmax.

- `get_fixed_m_constants(N)`: length-dependent recenter C(N) and safe bound
  that keep the f32 numerator/denominator from overflowing or flushing.
- `fixed_m_dtype_is_safe`: rejects fp16/fp8 (exponent range too small for the
  recentered exponent); bf16 and fp32 pass.
- V-magnitude gate: |V| <= DEFAULT_MAX_V_BOUND (256), otherwise online.
- Per-(head, Q-block) dispatch in the kernel; the Ulysses and R=1 ring callers
  compile a uniform (all-fixed) and a hybrid kernel and pick one with
  `lax.cond` on `all_fixed`. Both kernels are compiled; the extra compile
  time / program size has not been measured.
- On the R=1 `_ulysses_ring_custom_attention` path, when `batch > 1`,
  `devices_in_batch_sharding == 1`, and `(batch * num_kv_heads) % ulysses_shards == 0`,
  `batch` is folded into `heads` (`[1, B*H, S, D]` and `[1, B*Hkv, S, D]` before
  `shard_map`, with `k_mean` `[1, local_kv_heads, D]` computed inside `shard_map`)
  so `pallas_call` with scalar prefetch runs a single grid-batched kernel launch
  instead of a `vmap` serial loop.
- Virtual K-centering: the kernel receives k_mean and adds q . k_bar to the
  fixed bound in VMEM (an MXU dot_general on an (Hkv, 8, d) tiled k_mean,
  precision=HIGHEST because the gate has no slack at its edge). K itself is
  never re-written in HBM. k_mean is padded to the lane-padded head_dim after
  `_pad_data_for_flash`, so head_dim < 128 works (regression test at d=64).
  `save_residuals=True` with `k_mean is not None` raises `NotImplementedError`;
  `make_splash_mha` checks `bool(interpret)` so `interpret=False` keeps the
  custom VJP.

Explicit metadata contracts on the fixed-m ring path
----------------------------------------------------
1. Norm representation is declared, not inferred. The gate previously guessed
   whether `fixed_m_norms` were squared with `(qn.max() * mk.max()) < 1000.0`.
   Magnitude cannot answer that: legacy unsquared norms of (1000, 2) have a
   true bound of 2000, but read as squared give sqrt(2000) ~= 44.7, a ~45x
   under-estimate that admits fixed-m where it must fall back. Replaced by a
   REQUIRED `fixed_m_norms_squared` argument (no default; omitting it with
   fixed-m on raises). Tests use squared norms.

2. The V-safety predicate is required, not assumed. The V-magnitude and dtype
   verdict cannot be re-derived from one hop's Q/K, so the ring kernel raises
   unless `v_ok` is passed (`value=None` in the metadata helper gives
   v_ok = 0, fail-closed). `_ulysses_ring_custom_attention` (R > 1) computes it
   on the PRE-all-to-all shards -- dtype safety plus |V| <= 256 -- and reduces
   it with pmin over BOTH internal axes (ring, ulysses): before the a2a each
   device holds a sequence shard of every head, and the ring's
   accumulate-vs-LSE branch must be taken uniformly by every ppermute
   participant. Only Q and K go through the optimization_barrier that keeps
   XLA from duplicating their producer chains into the norm reductions; V is
   reduced outside it so the Q/K a2a does not wait on V's producer.

3. Norm shape is validated against per_q_block. Both gates compute
   `qn * mk[:, None]`, so a (num_heads,) array supplied with per_q_block=True
   broadcasts to (num_heads, num_heads) and pairs head j's query norm with head
   h's key norm; a sink head then inherits a small bound and the kernel
   evaluates exp2(large_logit - small_m) -> inf. The kernel now rejects the
   mismatch.

Tests
- custom_splash_fixed_m_test.py (37): kernel unit tests with relative-L2
  checks (including S < block_kv), k-norm (mk) shape and V-metadata fail-closed
  checks, `save_residuals + k_mean` and `interpret=False` VJP dispatch checks,
  `batch > 1` head-folding parity and HLO single-launch checks, integration
  through `_ulysses_attention` and the R=1 ring, and a U=2/R=2 check that a
  single V outlier on one pre-a2a device makes v_ok False on every device.
- ring_fixed_m_test.py (19): 8 backend-independent contract tests (including
  the required fixed_m_norms_squared and v_ok arguments), 8 2-device ring kernel
  tests (5 numerical tests vs an f32 reference, 1 test that declared-unsquared
  and squared norms give the same result, and 2 adversarial overflow tests
  verifying `v_ok=False` forces finite output while `v_ok=True` overflows fp32
  `pv`), and 3 tests showing the raw-K bound is unsound where the centered bound
  is not.
- CI: all 56 tests across `custom_splash_fixed_m_test.py` (37) and
  `ring_fixed_m_test.py` (19) run unconditionally in GitHub Actions (using Pallas
  interpret mode on 4 CPU host devices when not on TPU). The new
  `end_to_end/tpu/run_wan_stack_tests.sh` runs all of them on a TPU VM; later
  PRs in this stack add their test files to it.

Also pyink-reformats configuration_utils.py and wan_trainer_2_2.py, which
fail the Linter workflow on main (formatting only; same hunk as the
standalone chore PR so it rebases away once that merges).